import unittest

from psychsim.pwl import *
from psychsim.action import Action,ActionSet
from psychsim.world import World,stateKey
from psychsim.agent import Agent
from psychsim.reward import *

''' 
This example is used to test the dynamic turn taking
with an abstract self.director agent that acts every two turns
and can add or remove other self.agents from the turn queue
'''

QUEUE_LENGTH = 100.

NB_ROUNDS = 20

script = {3: ('remove','Pong'), 
          6: ('remove','Bob'),
          8: ('add-file','Pong'),
          12: ('add-pile','Bob'),
          15: ('remove', 'Pong')}

class AgentEjectionTestCase(unittest.TestCase):

    def setUp(self):
        #setup world, agent, and basics
        self.maxRounds= NB_ROUNDS
        self.world = World()
        self.ping = Agent('Ping')
        self.world.addAgent(self.ping)
        self.pong = Agent('Pong')
        self.world.addAgent(self.pong)
        self.bob = Agent('Bob')
        self.world.addAgent(self.bob)
        
        self.director = Agent('self.director')
        self.world.addAgent(self.director)
        
        self.agents = [self.ping, self.pong, self.bob]

        self.holdIndex = 10                             
        self.turn_list = ['self.director']+[a.name for a in self.agents]
        self.world.setOrder(self.turn_list, holdIndex = 10)        

    def addStates(self):
        #define states
        self.world.defineState(None,'round',int)
        self.world.setState(None,'round',0)
        self.world.defineState(None,'nb_self.agents',int)
        self.world.setState(None,'nb_self.agents',len(self.agents))

        for agent in self.agents:
            self.world.addDependency(turnKey(agent.name),stateKey(None,'nb_self.agents'))        
                    
    def addTermination(self):
        #terminate if maxRounds reached
        self.world.addTermination(makeTree({'if': thresholdRow(stateKey(None,'round'),self.maxRounds),
                                   True: True, 
                                   False: {'if': equalRow(stateKey(None,'nb_self.agents'),0),
                                           True: True, 
                                           False: False}}))

    def addActions(self):
        #add actions and appropriate dynamics
        for me in self.agents: 
            action =  me.addAction({'verb': 'do nothing'})
            tree = makeTree(incrementMatrix(stateKey(None,'round'),1)) 
            self.world.setDynamics('round',action,tree)
            tree = makeTree(setToFeatureMatrix(turnKey(me.name), stateKey(None,'nb_self.agents'), pct = 2., shift = -0.0)) 
            print tree
            self.world.setTurnDynamics(me.name, action, tree)
            #self.world.setDynamics(turnKey(me.name), action, tree)
            
        action = self.director.addAction({'verb': 'start'})
        for i in range(len(self.agents)):
            agent = self.agents[i]
            position = i*2+1
            tree = makeTree(setToConstantMatrix(turnKey(agent.name), position))
            self.world.setDynamics(turnKey(agent.name), action, tree)
        tree = makeTree(setToConstantMatrix(turnKey(self.director.name), 2))
        self.world.setTurnDynamics(self.director.name, action, tree)
        #self.world.setDynamics(turnKey(self.director.name), action, tree)
        self.director.setLegal(action, makeTree({'if': equalRow(stateKey(None,'round'),0),
                                   True:True,
                                   False:False}))
            
        action =  self.director.addAction({'verb': 'pass'})
        tree = makeTree(setToConstantMatrix(turnKey(self.director.name), 2)) 
        self.world.setDynamics(turnKey(self.director.name), action, tree)
        self.setScriptLegality(self.director, action)
        
        for ejected_agent in self.agents:
            action =  self.director.addAction({'verb': 'remove', 'object': ejected_agent.name})
            for moved_agent in self.agents:
                if moved_agent == ejected_agent:
                    tree = makeTree(setToConstantMatrix(turnKey(moved_agent.name),self.holdIndex))
                else:
                    tree = makeTree({'if': equalRow(turnKey(moved_agent.name),self.holdIndex),
                                     True: noChangeMatrix(turnKey(moved_agent.name)),
                                     False: {'if': greaterThanRow(turnKey(ejected_agent.name),turnKey(moved_agent.name)),
                                             True: incrementMatrix(turnKey(moved_agent.name),-1),#noChangeMatrix(turnKey(moved_agent.name)),
                                             False: setToFeatureMatrix(turnKey(moved_agent.name), turnKey(moved_agent.name), shift=-3.)}})
                    print str(tree)
                self.world.setDynamics(turnKey(moved_agent.name), action, tree)
            tree = makeTree(setToConstantMatrix(turnKey(self.director.name), 2)) 
            self.world.setDynamics(turnKey(self.director.name), action, tree)
            tree = makeTree(incrementMatrix(stateKey(None,'nb_self.agents'), -1.)) 
            self.world.setDynamics(stateKey(None,'nb_self.agents'), action, tree)
            self.setScriptLegality(self.director, action)
            
            
        for added_agent in self.agents:
            action =  self.director.addAction({'verb': 'add-file', 'object': added_agent.name})
            tree = makeTree(incrementMatrix(stateKey(None,'nb_self.agents'), 1.)) 
            self.world.setDynamics(stateKey(None,'nb_self.agents'), action, tree)
            tree = makeTree(setToFeatureMatrix(turnKey(added_agent.name),stateKey(None,'nb_self.agents'), pct=2., shift=-1.0))
            self.world.setDynamics(turnKey(added_agent.name), action, tree)
            tree = makeTree(setToConstantMatrix(turnKey(self.director.name), 2)) 
            self.world.setDynamics(turnKey(self.director.name), action, tree)
            self.setScriptLegality(self.director, action)
            
        for added_agent in self.agents:
            action =  self.director.addAction({'verb': 'add-pile', 'object': added_agent.name})
            tree = makeTree(incrementMatrix(stateKey(None,'nb_self.agents'), 1.)) 
            self.world.setDynamics(stateKey(None,'nb_self.agents'), action, tree)
            tree = makeTree(setToConstantMatrix(turnKey(added_agent.name),1))
            self.world.setDynamics(turnKey(added_agent.name), action, tree)
            for moved_agent in self.agents:
                if moved_agent != added_agent:
                    tree = makeTree({'if': equalRow(turnKey(moved_agent.name),-1),
                                     True: noChangeMatrix(turnKey(moved_agent.name)),
                                     False: setToFeatureMatrix(turnKey(moved_agent.name), turnKey(moved_agent.name), shift=1.)}) 
                    self.world.setDynamics(turnKey(moved_agent.name), action, tree)
            tree = makeTree(setToConstantMatrix(turnKey(self.director.name),2))
            self.world.setDynamics(turnKey(self.director.name), action, tree)
            self.setScriptLegality(self.director, action)

    def setScriptLegality(self, director_agent, action):
        if action['verb'] == 'pass':
            script_dict = True
            for k in script.keys():
                script_dict = {'if': equalRow(stateKey(None,'round'), k),
                                   True: False,
                                   False: script_dict}
            tree = makeTree({'if': equalRow(stateKey(None,'round'),0),
                                   True:False,
                                   False: script_dict})
        else:
            script_dict = False
            for k in script.keys():
                if script[k][0] == action['verb'] and script[k][1] == action['object']:
                    script_dict = {'if': equalRow(stateKey(None,'round'), k),
                                   True: True,
                                   False: script_dict}
            if action['verb'] == 'remove':
                tree = makeTree({'if': equalRow(turnKey(action['object']),self.holdIndex),#QUEUE_LENGTH),
                                 True:False,
                                 False: script_dict})
            else:
                tree = makeTree({'if': equalRow(turnKey(action['object']),self.holdIndex),#QUEUE_LENGTH),
                                 True: script_dict,
                                 False: False})
        director_agent.setLegal(action, tree)        

    def testSetup(self):
        self.addStates()
        self.addTermination()
        self.addActions()

    def testRun(self):
        self.testSetup()
        print str(self.world.state)
        while True:
            self.world.explain(self.world.step(),level=1)
            print str(self.world.state)
            if self.world.terminated():
                break 

if __name__ == "__main__":
    suite = unittest.TestSuite()
    for method in dir(AgentEjectionTestCase):
       if method.startswith("test"):
          suite.addTest(AgentEjectionTestCase(method))
    unittest.TextTestRunner().run(suite)
            